from dataclasses import dataclass
from typing import Any
from jinja2 import Template, FileSystemLoader, Environment


j2_loader = FileSystemLoader('./j2_files')
env = Environment(loader=j2_loader)

@dataclass
class PythonOperator:
    operator_name: str = "operator"
    operator_type: str = "PythonOperator"
    task_id: str = ""
    python_callable: str = "python_function"
    op_kwargs: list[str] = None
    dag: str = "dag"

    def create_python_operator(self):
        # 生成 PythonOperator 字符串
        PythonOperator_j2 = env.get_template('./PythonOperator.j2')
        result = PythonOperator_j2.render(
            operator_name=self.operator_name,
            operator_type=self.operator_type,
            task_id=self.task_id,
            python_callable=self.python_callable,
            op_kwargs=self.op_kwargs,
            dag=self.dag,
        )
        return result

@dataclass
class SparkKubernetesOperator:
    operator_name: str = "operator"
    operator_type: str = "SparkKubernetesOperator"
    task_id: str = ""
    namespace: str = "default"
    application_file: str = "spark_job_template.yaml"
    image: str = "spark-pandas:3.5.3"
    code_path: str = "local:///home/count_schools.py"
    random_name_suffix: bool = True
    dag: str = "dag"

    def create_spark_kubernetes_operator(self):
        # 生成 SparkKubernetesOperator 字符串
        SparkKubernetesOperator_j2 = env.get_template('./SparkKubernetesOperator.j2')
        result = SparkKubernetesOperator_j2.render(
            operator_name=self.operator_name,
            operator_type=self.operator_type,
            task_id=self.task_id,
            namespace=self.namespace,
            application_file=self.application_file,
            image=self.image,
            code_path=self.code_path,
            random_name_suffix=self.random_name_suffix,
            dag=self.dag,
        )
        return result